# -*- coding: utf-8 -*-
"""
Created on Sat Jan 18 13:26:56 2020
本程序参考EAST的dataset.py改写
@author: ikaros
"""
'''
本文件的函数：
get_rotate_mat(theta):                          计算旋转矩阵
rotate_vertices(vertices, theta, anchor=None):  旋转目标区域的坐标
rotate_img(img, vertices, angle_range=r_angle): 旋转图片
extract_vertices(lines):                        从标注文件里提取坐标值跟标签值
check_verticles(imgdir,labeldir,out_dir):       检查旋转后图片的坐标是否正确，最后也print出错误的文件名，需手动删除；还需到checkdir下检查一下图片效果。    
'''

import numpy as np
import cv2
from PIL import Image
import math
import os #路径有中文时可能会报一些奇怪的错

'''---------4个旋转细节函数------------'''
def get_rotate_mat(theta):
	'''positive theta value means rotate clockwise'''
	return np.array([[math.cos(theta), -math.sin(theta)], [math.sin(theta), math.cos(theta)]])


def rotate_vertices(vertices, theta, anchor=None):
	'''rotate vertices around anchor
	Input:	
		vertices: vertices of text region <numpy.ndarray, (8,)>
		theta   : angle in radian measure
		anchor  : fixed position during rotation
	Output:
		rotated vertices <numpy.ndarray, (8,)>
	'''
	v = vertices.reshape((4,2)).T
	if anchor is None:
		anchor = v[:,:1]
	rotate_mat = get_rotate_mat(theta)
	res = np.dot(rotate_mat, v - anchor)
	return (res + anchor).T.reshape(-1)

#def rotate_img(img, vertices, angle_range=r_angle):
def rotate_img(img, vertices, angle_range):
    '''
	Input:
		img         : PIL Image
		vertices    : vertices of text regions <numpy.ndarray, (n,8)>
		angle_range : rotate range
	Output:
		img         : rotated PIL Image
		new_vertices: rotated vertices
    '''

    center_x = (img.width - 1) / 2
    center_y = (img.height - 1) / 2
#    angle = angle_range * (np.random.rand() * 2 - 1)
    angle = angle_range
    img = img.rotate(angle, Image.BILINEAR)
    new_vertices = np.zeros(vertices.shape)
    for i, vertice in enumerate(vertices):
        new_vertices[i,:] = rotate_vertices(vertice, -angle / 180 * math.pi, np.array([[center_x],[center_y]]))
    return img, new_vertices

def extract_vertices(lines):
	'''extract vertices info from txt lines
	Input:
		lines   : 标签文件里的所有内容
	Output:
		vertices: vertices of text regions <numpy.ndarray, (n,8)>
		labels  : 1->valid, 0->ignore, <numpy.ndarray, (n,)>
	'''
	labels = []
	vertices = []
	for line in lines:
		vertices.append(list(map(int,list(map(float,line.rstrip('\n').rstrip().lstrip('\ufeff').split(',')[:8])))))
		label = 0 if '###' in line else 1
		labels.append(label)
	return np.array(vertices), np.array(labels)
'''-------------------------------------------------------------------------------'''

'''----------------------------------------------------------------------------------------------------------------'''
def label_file_delete(img_dir, label_dir):
    '''删除多余的标签文件'''
    #imgnames = np.array(os.listdir(img_dir))
    labelnames = np.array(os.listdir(label_dir))
    num = 0
    for labelname in labelnames:
        labelpath = label_dir + labelname
        imgpath = img_dir + labelname.replace('gt_','').split('.')[0] + '.jpg'
        
        if not os.path.exists(imgpath):
            num += 1
            #print(labelname)
            os.remove(labelpath)
    print('deletion number : ', num)
    

def aug_rotate_img_and_label(o_img_root, o_label_root, save_img_root, save_label_root, r_angle):
    """旋转图片----->数据增强"""
    
    o_img_names = np.array(os.listdir(o_img_root))        #读取所有图片的地址和名字 #路径不能有中文
    o_label_names = np.array(os.listdir(o_label_root))    #读取所有图片的地址和名字 #路径不能有中文
    #print('o_img shape: ',o_img_names.shape,'  ; o_label shape: ', o_label_names.shape)
    #print('o_img size: ',o_img_names.size,'  ; o_label size: ', o_label_names.size)
    
    if not os.path.exists(save_img_root):                 #创建保存目录
        os.makedirs(save_img_root)
    if not os.path.exists(save_label_root):
        os.makedirs(save_label_root)
    
    symbol = '(r' + str(r_angle) + ')'
   
    for file in o_img_names:     
        print(type(file))                        
        
        o_img_name = o_img_root + file                                        #图片所在的真实地址
        new_img_name = save_img_root + file.split('.')[0] + symbol + '.jpg'    #保存图片的绝对地址
        o_label_name = o_label_root + 'gt_' + file.replace('.jpg','.txt')
        new_label_name = save_label_root + 'gt_' + file.split('.')[0] + symbol + '.txt'
        
        with open(o_label_name, 'r') as f:
            lines = f.readlines()
    
        img = cv2.imread(o_img_name)                  #读取图片
        if img is None:                               #跳过有问题的图片
            continue
        
        img2 = Image.open(o_img_name)                 #使用PIL打开图片
        vertices, labels = extract_vertices(lines)    #提取图片的坐标与标签，默认所有标签为正样本
        img3, new_vertices = rotate_img(img2,vertices,r_angle)
        img3.save(new_img_name)                       #保存旋转后的图片
        
        '''把旋转后坐标取整并转化为字符'''
        newlines = []
        for v1 in new_vertices:
            if v1 is None:
                continue
            x1 = str(int(round(v1[0])))               #四舍五入取整数，一般像素点的坐标为整数
            y1 = str(int(round(v1[1])))
            x2 = str(int(round(v1[2])))
            y2 = str(int(round(v1[3])))
            x3 = str(int(round(v1[4])))
            y3 = str(int(round(v1[5])))
            x4 = str(int(round(v1[6])))
            y4 = str(int(round(v1[7])))
            
            newline = x1 + ',' + y1 + ',' + x2 + ',' + y2 + ','+ x3 + ','+ y3 + ','+ x4 + ','+ y4 + ',' + 'A'+'\n' #A表示正样本，'###'表示负样本  
            print('newline: ',newline)
            newlines.append(newline)
        
        with open(new_label_name, "a") as f1:         #把坐标和标签内容写入txt文件
            for ns in newlines:
                f1.write(ns)
        #break
 

def check_verticles(imgdir,labeldir,out_dir):
    '''
    旋转完后图片后，
    用来检查坐标是否超出图片的大小范围，
    以及检查是否有的图片丢失标签文件。
    '''
    #若标签数多于图片数，
    #可以遍历标签名，把标签名转化图片名，
    #使用if not (img_names == imgname).any() 进行判断

    if not os.path.exists(out_dir):                 #创建保存目录
        os.makedirs(out_dir)
    
    imgnames = np.array(os.listdir(imgdir))
    

    label_names = np.array(os.listdir(labeldir))         #获取所有标签文件名
    num=0
    for imgname in imgnames:
        imgpath = imgdir + imgname
        img = cv2.imread(imgpath)
        w = np.array(img).shape[1]                       #宽高用来判断坐标是否超出图片范围，从而导致训练出错
        h = np.array(img).shape[0]
        num += 1                                         #用来计数
        labelname = 'gt_' + imgname.split('.')[0] + '.txt'  #把图片名转换成标签名
        #labelname = file.split('.')[0] + '.txt'  #把图片名转换成标签名
        #print(labelname)
        
        '''#判断标签数是否少于图片数'''
        if not (label_names == labelname).any():         
            print('no label-->labelname:',labelname,' False')
        
        label2 = labeldir + labelname
        if (label_names == labelname).any():
            with open(label2, "r") as f1:                #读取标签文件的坐标和标签
                lines = f1.readlines()
            vertices, label3 = extract_vertices(lines)   #提取坐标
            lines2 = np.array(lines)
            if lines2.shape[0] == 0:                     #检查标签文件内容是否为空
                print('labelname--','shape wrong:',labelname,':',lines2.shape[0])
            '''把坐标取整'''
            for v1 in vertices:
                if v1 is None:
                    print('None')
                    break
                x1 = round(v1[0])
                y1 = round(v1[1])
                x2 = round(v1[2])
                y2 = round(v1[3])
                x3 = round(v1[4])
                y3 = round(v1[5])
                x4 = round(v1[6])
                y4 = round(v1[7])
                if x1<0 or x2<0 or x3<0 or x4<0 or y1<0 or y2<0 or y3<0 or y4<0:
                    print('v_wrong labelname:',labelname,' False')
                if x1>w or x2>w or x3>w or x4>w:
                    print('v_wrong labelname:',labelname,' False')
                if y1>h or y2>h or y3>h or y4>h:
                    print('v_wrong labelname:',labelname,' False')
                
                '''用直线连接 绘制四边形'''
                '''【注意box的格式】'''
                box = np.array([[x1,y1],[x2,y2],[x3,y3],[x4,y4]],np.int32)
                box = box.reshape(-1,1,2)
                cv2.polylines(img, [box], True,color=(0, 255, 0),thickness=5)
            #img = cv2.resize(img,(500,500))
            outname = out_dir + 'gt_' + imgname.split('.')[0] + '.jpg'
            cv2.imwrite(outname,img)
'''----------------------------------------------------------------------------------------------------------------'''        


'''----------------------------------main 函数----------------------------------------'''

if __name__ == '__main__':
    root_dir = 'D:/temp_data/code/dataset_tool/'
    o_img_root= root_dir + 'origin_data/image/'                   #原始图片地址
    o_label_root = root_dir + 'origin_data/label/'              #原始标签地址
    save_img_root = root_dir + 'done_data/rotate_img/'          #新图片保存地址
    save_label_root = root_dir + 'done_data/rotate_label/'      #新标签保存地址
    checkdir = root_dir + 'check/'                   #验证图片是否正确的函数输出地址
    r_angle = -4                                     #图片旋转的角度
    
    #aug_rotate_img_and_label(o_img_root, o_label_root, save_img_root, save_label_root, r_angle) #旋转图片
    check_verticles(save_img_root,save_label_root,checkdir)                                      #检查图片
















